import os
import sys
import threading
import time
from datetime import datetime
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor

import cv2
import numpy as np
from PIL import Image
from PySide6.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, 
                             QHBoxLayout, QPushButton, QLabel, QTextEdit, 
                             QProgressBar, QTreeWidget, QTreeWidgetItem,
                             QFileDialog, QSpinBox, QDoubleSpinBox, QGroupBox,
                             QSplitter, QMessageBox, QHeaderView, QLineEdit)
from PySide6.QtCore import Qt, QThread, Signal
from PySide6.QtGui import QFont, QPalette, QColor

# ============================ 核心检测算法 ============================

def get_image_files(directory):
    """获取目录下所有图片文件"""
    image_extensions = {'.png', '.jpg', '.jpeg', '.bmp', '.tiff', '.webp'}
    image_files = []
    
    for root, dirs, files in os.walk(directory):
        for file in files:
            ext = os.path.splitext(file)[1].lower()
            if ext in image_extensions:
                image_files.append(os.path.join(root, file))
    
    return image_files

def calculate_hash(image_path):
    """计算图像的差异哈希（dHash）"""
    try:
        img = Image.open(image_path)
        img = img.convert('L').resize((9, 8), Image.LANCZOS)
        pixels = list(img.getdata())
        
        diff = []
        for i in range(8):
            for j in range(8):
                if pixels[i*9 + j] > pixels[i*9 + j+1]:
                    diff.append(1)
                else:
                    diff.append(0)
        
        hash_val = 0
        for i, bit in enumerate(diff):
            if bit:
                hash_val |= 1 << i
                
        return hash_val
    except Exception as e:
        print(f"计算图像哈希时出错 {image_path}: {e}")
        return None

def hamming_distance(hash1, hash2):
    """计算两个哈希值的汉明距离"""
    if hash1 is None or hash2 is None:
        return 64
    return bin(hash1 ^ hash2).count('1')

def similarity_percentage(hash1, hash2):
    """基于汉明距离计算相似度百分比"""
    distance = hamming_distance(hash1, hash2)
    return (1 - distance / 64) * 100

def group_similar_images(similar_pairs):
    """将相似图片对分组为相似的图片组"""
    parent = {}
    
    def find(x):
        if x not in parent:
            parent[x] = x
        if parent[x] != x:
            parent[x] = find(parent[x])
        return parent[x]
    
    def union(x, y):
        root_x, root_y = find(x), find(y)
        if root_x != root_y:
            parent[root_y] = root_x
    
    for path1, path2, _ in similar_pairs:
        union(path1, path2)
    
    groups = defaultdict(list)
    for path in parent:
        root = find(path)
        groups[root].append(path)
    
    result_groups = {}
    group_id = 1
    for root, paths in groups.items():
        if len(paths) > 1:
            result_groups[group_id] = sorted(paths)
            group_id += 1
    
    return result_groups

# ============================ 后台检测线程 ============================

class DetectionThread(QThread):
    """检测线程，避免GUI冻结"""
    progress_updated = Signal(int, int, str)  # 当前值, 最大值, 状态文本
    log_message = Signal(str)  # 日志消息
    result_ready = Signal(dict)  # 检测结果
    finished_success = Signal()  # 完成信号
    
    def __init__(self, directory, threshold=90, max_workers=10):
        super().__init__()
        self.directory = directory
        self.threshold = threshold
        self.max_workers = max_workers
        self._is_running = True
    
    def stop_detection(self):
        """停止检测"""
        self._is_running = False
    
    def run(self):
        """执行检测任务"""
        try:
            self.log_message.emit(f"开始扫描目录: {self.directory}")
            image_files = get_image_files(self.directory)
            total_images = len(image_files)
            self.log_message.emit(f"找到 {total_images} 个图片文件")
            
            if total_images < 2:
                self.log_message.emit("需要至少2个图片文件进行比较")
                self.result_ready.emit({})
                return
            
            # 阶段1：多线程计算所有图片的哈希值
            self.log_message.emit("开始计算图片哈希值...")
            hashes_dict = {}
            processed_count = 0
            
            # 将图片分成批次进行并行处理
            batch_size = max(1, total_images // self.max_workers)
            batches = [image_files[i:i+batch_size] for i in range(0, len(image_files), batch_size)]
            
            with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
                futures = []
                for batch in batches:
                    futures.append(executor.submit(self.compute_hashes_batch, batch))
                
                for i, future in enumerate(futures):
                    if not self._is_running:
                        break
                    batch_results = future.result()
                    for path, hash_val in batch_results:
                        if hash_val is not None:
                            hashes_dict[path] = hash_val
                        processed_count += 1
                    
                    # 更新进度（第一阶段占50%）
                    progress = int((processed_count / total_images) * 50)
                    self.progress_updated.emit(progress, 100, f"哈希计算: {processed_count}/{total_images}")
            
            if not self._is_running:
                return
                
            self.log_message.emit(f"成功计算 {len(hashes_dict)} 个图片的哈希值")
            
            if len(hashes_dict) < 2:
                self.log_message.emit("可用的图片数量不足，无法进行比较")
                self.result_ready.emit({})
                return
            
            # 阶段2：多线程比较图片相似度
            self.log_message.emit("开始比较图片相似度...")
            similar_pairs = []
            hash_items = list(hashes_dict.items())
            total_comparisons = len(hash_items) * (len(hash_items) - 1) // 2
            completed_comparisons = 0
            
            # 使用线程池进行并行比较
            def compare_chunk(start_idx, end_idx):
                chunk_results = []
                for i in range(start_idx, end_idx):
                    if not self._is_running:
                        break
                    path1, hash1 = hash_items[i]
                    for j in range(i + 1, len(hash_items)):
                        if not self._is_running:
                            break
                        path2, hash2 = hash_items[j]
                        similarity = similarity_percentage(hash1, hash2)
                        if similarity >= self.threshold:
                            chunk_results.append((path1, path2, similarity))
                        
                        completed_comparisons += 1
                        if completed_comparisons % 1000 == 0:
                            progress = 50 + int((completed_comparisons / total_comparisons) * 50)
                            self.progress_updated.emit(min(progress, 99), 100, f"比较进度: {completed_comparisons}/{total_comparisons}")
                
                return chunk_results
            
            chunk_size = max(1, len(hash_items) // (self.max_workers * 2))
            with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
                futures = []
                for start_idx in range(0, len(hash_items), chunk_size):
                    end_idx = min(start_idx + chunk_size, len(hash_items))
                    futures.append(executor.submit(compare_chunk, start_idx, end_idx))
                
                for future in futures:
                    if not self._is_running:
                        break
                    chunk_results = future.result()
                    similar_pairs.extend(chunk_results)
            
            if not self._is_running:
                return
                
            self.log_message.emit(f"完成比较，找到 {len(similar_pairs)} 个相似图片对")
            
            # 分组相似图片
            similar_groups = group_similar_images(similar_pairs)
            self.log_message.emit(f"分组完成，共 {len(similar_groups)} 组相似图片")
            
            # 确保进度条完成
            self.progress_updated.emit(100, 100, "检测完成")
            
            self.result_ready.emit(similar_groups)
            self.finished_success.emit()
            
        except Exception as e:
            self.log_message.emit(f"检测过程中出现错误: {str(e)}")
            self.result_ready.emit({})
    
    def compute_hashes_batch(self, image_paths):
        """批量计算图像哈希值"""
        results = []
        for path in image_paths:
            hash_val = calculate_hash(path)
            results.append((path, hash_val))
        return results

# ============================ 主GUI界面 ============================

class SimilarityDetectorGUI(QMainWindow):
    """图片相似度检测GUI主窗口"""
    
    def __init__(self):
        super().__init__()
        self.detection_thread = None
        self.current_results = {}
        self.init_ui()
        
    def init_ui(self):
        """初始化用户界面"""
        self.setWindowTitle("图片相似度检测系统")
        self.setGeometry(100, 100, 1200, 800)
        
        # 创建中央部件
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        
        # 主布局
        main_layout = QVBoxLayout(central_widget)
        
        # 创建标题
        title_label = QLabel("图片相似度检测系统")
        title_font = QFont()
        title_font.setPointSize(16)
        title_font.setBold(True)
        title_label.setFont(title_font)
        title_label.setAlignment(Qt.AlignCenter)
        title_label.setStyleSheet("QLabel { color: #2c3e50; margin: 10px; }")
        main_layout.addWidget(title_label)
        
        # 创建分割器
        splitter = QSplitter(Qt.Vertical)
        main_layout.addWidget(splitter)
        
        # 上部区域：控制和参数设置
        top_widget = QWidget()
        top_layout = QVBoxLayout(top_widget)
        
        # 目录选择区域
        dir_group = QGroupBox("目录选择")
        dir_layout = QHBoxLayout(dir_group)
        
        self.dir_label = QLabel("未选择目录")
        self.dir_label.setStyleSheet("QLabel { padding: 8px; border: 1px solid #bdc3c7; border-radius: 4px; background: #ecf0f1; }")
        self.dir_label.setMinimumHeight(35)
        
        self.select_dir_btn = QPushButton("选择图片目录")
        self.select_dir_btn.setStyleSheet("QPushButton { padding: 8px 15px; background: #3498db; color: white; border: none; border-radius: 4px; }"
                                         "QPushButton:hover { background: #2980b9; }")
        self.select_dir_btn.clicked.connect(self.select_directory)
        
        dir_layout.addWidget(QLabel("目标目录:"))
        dir_layout.addWidget(self.dir_label, 1)
        dir_layout.addWidget(self.select_dir_btn)
        
        # 参数设置区域
        param_group = QGroupBox("检测参数")
        param_layout = QHBoxLayout(param_group)
        
        param_layout.addWidget(QLabel("相似度阈值 (%):"))
        self.threshold_spin = QDoubleSpinBox()
        self.threshold_spin.setRange(0, 100)
        self.threshold_spin.setValue(98.0)  # 默认值设为98%
        self.threshold_spin.setDecimals(2)
        self.threshold_spin.setSingleStep(1)
        self.threshold_spin.setSuffix("%")
        
        param_layout.addWidget(self.threshold_spin)
        param_layout.addWidget(QLabel("线程数:"))
        self.workers_spin = QSpinBox()
        self.workers_spin.setRange(1, 20)
        self.workers_spin.setValue(10)
        
        param_layout.addWidget(self.workers_spin)
        param_layout.addStretch(1)
        
        # 控制按钮区域
        control_layout = QHBoxLayout()
        
        self.start_btn = QPushButton("开始检测")
        self.start_btn.setStyleSheet("QPushButton { padding: 10px 20px; background: #27ae60; color: white; border: none; border-radius: 4px; font-weight: bold; }"
                                    "QPushButton:hover { background: #229954; }"
                                    "QPushButton:disabled { background: #bdc3c7; }")
        self.start_btn.clicked.connect(self.start_detection)
        self.start_btn.setEnabled(False)
        
        self.stop_btn = QPushButton("停止检测")
        self.stop_btn.setStyleSheet("QPushButton { padding: 10px 20px; background: #e74c3c; color: white; border: none; border-radius: 4px; font-weight: bold; }"
                                   "QPushButton:hover { background: #c0392b; }"
                                   "QPushButton:disabled { background: #bdc3c7; }")
        self.stop_btn.clicked.connect(self.stop_detection)
        self.stop_btn.setEnabled(False)
        
        self.export_btn = QPushButton("导出结果")
        self.export_btn.setStyleSheet("QPushButton { padding: 10px 20px; background: #f39c12; color: white; border: none; border-radius: 4px; font-weight: bold; }"
                                     "QPushButton:hover { background: #d35400; }"
                                     "QPushButton:disabled { background: #bdc3c7; }")
        self.export_btn.clicked.connect(self.export_results)
        self.export_btn.setEnabled(False)
        
        control_layout.addWidget(self.start_btn)
        control_layout.addWidget(self.stop_btn)
        control_layout.addWidget(self.export_btn)
        control_layout.addStretch(1)
        
        # 进度条
        self.progress_bar = QProgressBar()
        self.progress_bar.setRange(0, 100)
        self.progress_bar.setVisible(False)
        self.progress_bar.setStyleSheet("QProgressBar { border: 1px solid #bdc3c7; border-radius: 4px; text-align: center; }"
                                       "QProgressBar::chunk { background: #3498db; }")
        
        self.status_label = QLabel("准备就绪")
        self.status_label.setAlignment(Qt.AlignCenter)
        self.status_label.setStyleSheet("QLabel { color: #7f8c8d; padding: 5px; }")
        
        # 添加到上部布局
        top_layout.addWidget(dir_group)
        top_layout.addWidget(param_group)
        top_layout.addLayout(control_layout)
        top_layout.addWidget(self.progress_bar)
        top_layout.addWidget(self.status_label)
        
        # 下部区域：日志和结果展示
        bottom_widget = QWidget()
        bottom_layout = QHBoxLayout(bottom_widget)
        
        # 日志区域
        log_group = QGroupBox("检测日志")
        log_layout = QVBoxLayout(log_group)
        
        self.log_text = QTextEdit()
        self.log_text.setReadOnly(True)
        self.log_text.setStyleSheet("QTextEdit { background: #f8f9fa; border: 1px solid #dee2e6; border-radius: 4px; }")
        log_layout.addWidget(self.log_text)
        
        # 结果区域
        result_group = QGroupBox("检测结果")
        result_layout = QVBoxLayout(result_group)
        
        self.result_tree = QTreeWidget()
        self.result_tree.setHeaderLabels(["相似图片组", "图片路径", "文件大小"])
        self.result_tree.header().setSectionResizeMode(QHeaderView.Interactive)
        self.result_tree.setStyleSheet("QTreeWidget { background: #f8f9fa; border: 1px solid #dee2e6; border-radius: 4px; }")
        result_layout.addWidget(self.result_tree)
        
        # 左右分割
        bottom_layout.addWidget(log_group, 1)
        bottom_layout.addWidget(result_group, 1)
        
        # 添加到分割器
        splitter.addWidget(top_widget)
        splitter.addWidget(bottom_widget)
        splitter.setSizes([300, 500])
        
        # 初始化日志
        self.log_message("欢迎使用图片相似度检测工具")
        self.log_message("请先选择包含图片的目录")
        
    def select_directory(self):
        """选择图片目录"""
        directory = QFileDialog.getExistingDirectory(
            self, 
            "选择图片目录", 
            "",
            QFileDialog.Option.ShowDirsOnly
        )
        
        if directory:
            self.directory = directory
            self.dir_label.setText(directory)
            self.start_btn.setEnabled(True)
            self.log_message(f"已选择目录: {directory}")
            
            # 预览目录中的图片数量
            try:
                image_files = get_image_files(directory)
                self.log_message(f"目录中包含 {len(image_files)} 个图片文件")
            except Exception as e:
                self.log_message(f"扫描目录时出错: {str(e)}")
    
    def start_detection(self):
        """开始检测过程"""
        if not hasattr(self, 'directory') or not self.directory:
            QMessageBox.warning(self, "警告", "请先选择图片目录")
            return
        
        # 禁用开始按钮，启用停止按钮
        self.start_btn.setEnabled(False)
        self.stop_btn.setEnabled(True)
        self.export_btn.setEnabled(False)
        
        # 清空之前的结果
        self.result_tree.clear()
        
        # 显示进度条
        self.progress_bar.setVisible(True)
        self.progress_bar.setValue(0)
        
        # 获取参数
        threshold = self.threshold_spin.value()
        max_workers = self.workers_spin.value()
        
        self.log_message("开始图片相似度检测...")
        self.log_message(f"参数 - 相似度阈值: {threshold}%, 线程数: {max_workers}")
        
        # 创建并启动检测线程
        self.detection_thread = DetectionThread(self.directory, threshold, max_workers)
        self.detection_thread.progress_updated.connect(self.update_progress)
        self.detection_thread.log_message.connect(self.log_message)
        self.detection_thread.result_ready.connect(self.detection_finished)
        self.detection_thread.finished_success.connect(self.detection_success)
        self.detection_thread.start()
    
    def stop_detection(self):
        """停止检测过程"""
        if self.detection_thread and self.detection_thread.isRunning():
            self.detection_thread.stop_detection()
            self.detection_thread.wait()
            self.log_message("检测已停止")
        
        self.start_btn.setEnabled(True)
        self.stop_btn.setEnabled(False)
        self.progress_bar.setVisible(False)
        self.status_label.setText("检测已停止")
    
    def update_progress(self, value, maximum, status):
        """更新进度条和状态"""
        self.progress_bar.setMaximum(maximum)
        self.progress_bar.setValue(value)
        self.status_label.setText(status)
    
    def detection_finished(self, results):
        """检测完成处理"""
        self.current_results = results
        
        # 修复逻辑：正确显示检测结果
        if results:
            self.log_message(f"检测完成！共找到 {len(results)} 组相似图片")
            self.display_results(results)
        else:
            self.log_message("检测完成，未找到相似度超过阈值的图片")
    
    def detection_success(self):
        """检测成功完成"""
        self.start_btn.setEnabled(True)
        self.stop_btn.setEnabled(False)
        self.export_btn.setEnabled(len(self.current_results) > 0)
        self.progress_bar.setVisible(False)
        self.status_label.setText("检测完成")
    
    def display_results(self, results):
        """在树形控件中显示结果"""
        self.result_tree.clear()
        
        for group_id, image_paths in results.items():
            group_item = QTreeWidgetItem(self.result_tree)
            group_item.setText(0, f"第 {group_id} 组 ({len(image_paths)} 张图片)")
            group_item.setExpanded(True)
            
            for path in image_paths:
                file_item = QTreeWidgetItem(group_item)
                file_item.setText(0, "")  # 空字符串，使路径从第二列开始
                file_item.setText(1, os.path.basename(path))
                file_item.setText(2, self.get_file_size(path))
                
                # 完整路径作为工具提示
                file_item.setToolTip(1, path)
        
        # 调整列宽
        for i in range(self.result_tree.columnCount()):
            self.result_tree.resizeColumnToContents(i)
    
    def get_file_size(self, path):
        """获取文件大小并格式化"""
        try:
            size = os.path.getsize(path)
            if size < 1024:
                return f"{size} B"
            elif size < 1024 * 1024:
                return f"{size/1024:.1f} KB"
            else:
                return f"{size/(1024 * 1024):.1f} MB"
        except:
            return "未知"
    
    def export_results(self):
        """导出检测结果到文件"""
        if not self.current_results:
            QMessageBox.information(self, "提示", "没有可导出的结果")
            return
        
        file_path, _ = QFileDialog.getSaveFileName(
            self,
            "导出结果",
            f"相似度检测结果_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt",
            "文本文件 (*.txt);;所有文件 (*)"
        )
        
        if file_path:
            try:
                with open(file_path, 'w', encoding='utf-8') as f:
                    f.write("图片相似度检测报告\n")
                    f.write("=" * 50 + "\n")
                    f.write(f"生成时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
                    f.write(f"扫描目录: {self.directory}\n")
                    f.write(f"相似度阈值: {self.threshold_spin.value()}%\n")
                    f.write(f"找到 {len(self.current_results)} 组相似图片\n\n")
                    
                    for group_id, image_paths in self.current_results.items():
                        f.write(f"第 {group_id} 组相似图片 ({len(image_paths)} 张):\n")
                        for path in image_paths:
                            rel_path = os.path.relpath(path, self.directory)
                            file_size = self.get_file_size(path)
                            f.write(f"  - {rel_path} ({file_size})\n")
                        f.write("\n")
                
                self.log_message(f"结果已导出到: {file_path}")
                QMessageBox.information(self, "成功", f"结果已导出到:\n{file_path}")
            except Exception as e:
                self.log_message(f"导出结果时出错: {str(e)}")
                QMessageBox.critical(self, "错误", f"导出失败: {str(e)}")
    
    def log_message(self, message):
        """添加日志消息"""
        timestamp = datetime.now().strftime("%H:%M:%S")
        self.log_text.append(f"[{timestamp}] {message}")
        # 自动滚动到底部
        scrollbar = self.log_text.verticalScrollBar()
        scrollbar.setValue(scrollbar.maximum())
    
    def closeEvent(self, event):
        """窗口关闭事件"""
        if hasattr(self, 'detection_thread') and self.detection_thread and self.detection_thread.isRunning():
            reply = QMessageBox.question(
                self, 
                "确认退出", 
                "检测正在进行中，确定要退出吗？",
                QMessageBox.Yes | QMessageBox.No,
                QMessageBox.No
            )
            
            if reply == QMessageBox.Yes:
                self.detection_thread.stop_detection()
                self.detection_thread.wait()
                event.accept()
            else:
                event.ignore()
        else:
            event.accept()

def main():
    """主函数"""
    app = QApplication(sys.argv)
    
    # 设置应用程序样式
    app.setStyle('Fusion')
    
    # 创建并显示主窗口
    window = SimilarityDetectorGUI()
    window.show()
    
    # 运行应用程序
    sys.exit(app.exec())

if __name__ == "__main__":
    main()